from typing import List, Dict

from tree_rag.dataclasses.api import Knowledge
from tree_rag.session import Message

from . import DEFAULT_MODEL_NAME, Prompt, client

SYSTEM_PROMPT = """{role_prompt}
你的任务是根据检索到的知识对用户的问题进行回答，回答要精准，简短，准确
要求：
- 只根据给定的上下文知识进行回答，无法从上下文知识中得到答案的，请如实告知用户。
- 要求精准定位用户问题，答案力求简洁，完整，明了。尽量条目清晰，要点分行分段。
- 所有的答案都要尽量引用检索到的知识，不要编造和自由发挥。
- 不要漏掉和问题相关的内容
- 不要自由发挥，不要自由发挥，不要自由发挥，重要的事情说三遍。回答的每句话要有引用依据。
"""

USER_PROMPT = """
检索到的知识：
{context}

{exclusion_block}

对话记录：
{history}

请根据检索到的知识对用户的问题进行回答，回答要精准，简短，准确，不要自由发挥，不要漏掉和问题相关的内容。

user: {user_input}
"""


SYSTEM_PROMPT_USELLM_FOR_UNKNOWN = """{role_prompt}
你的任务是根据检索到的知识对用户的问题进行回答，回答要精准，简短，准确
要求：
- 要求精准定位用户问题，答案力求简洁，完整，明了。尽量条目清晰，要点分行分段。
- 所有的答案都要尽量引用检索到的知识，如果检索不到答案，请自由发挥。
- 不要漏掉和问题相关的内容
"""

USER_PROMPT_USELLM_FOR_UNKNOWN = """
检索到的知识：
{context}

{exclusion_block}

对话记录：
{history}

请根据检索到的知识对用户的问题进行回答，回答要精准，简短，准确，如果检索到的知识不足以回答问题，请自由发挥，不要漏掉和问题相关的内容。

user: {user_input}
"""


def build_prompt(
    knowledges: List[Knowledge],
    histories: List[Message],
    role_prompt: str,
    use_llm_for_unknown=False,
    exclusion: List[Dict] = None,
) -> Prompt:
    context_str = "\n\n".join([item.to_str_context() for item in knowledges])
    history = "\n".join([item.to_str() for item in histories[:-1]])
    
    # Build exclusion block
    exclusion_lines: List[str] = []
    if exclusion:
        exclusion_lines.append("另外，下列类别也无需使用固定拒答话术，请按类别自然回复：")
        for item in exclusion:
            category = item.get("category") if isinstance(item, dict) else getattr(item, "category", None)
            examples = item.get("examples") if isinstance(item, dict) else getattr(item, "examples", [])
            definition = item.get("definition") if isinstance(item, dict) else getattr(item, "definition", None)
            words = item.get("words") if isinstance(item, dict) else getattr(item, "words", [])
            
            if category:
                # Build a more natural description for each category
                category_desc = f"**{category}**"
                
                if definition:
                    # Clean up definition formatting and add it naturally
                    clean_definition = definition.strip().replace('\n', ' ')
                    category_desc += f"：{clean_definition}"
                
                if examples and len(examples) > 0:
                    examples_str = "、".join(examples)
                    category_desc += f" 例如：{examples_str}。"
                elif definition:
                    category_desc += "。"
                
                if words and len(words) > 0:
                    # Use the first response template as the main one
                    main_response = words[0] if isinstance(words[0], str) else str(words[0])
                    category_desc += f" 建议回复：{main_response}"
                
                exclusion_lines.append(f"- {category_desc}")
    exclusion_block = ("\n" + "\n".join(exclusion_lines)) if exclusion_lines else ""
    
    if not use_llm_for_unknown:
        user_prompt = USER_PROMPT.format(
            context=context_str,
            history=history,
            exclusion_block=exclusion_block,
            user_input=histories[-1].content,
        )
        system_prompt = SYSTEM_PROMPT.format(role_prompt=role_prompt)
    else:
        user_prompt = USER_PROMPT_USELLM_FOR_UNKNOWN.format(
            context=context_str,
            history=history,
            exclusion_block=exclusion_block,
            user_input=histories[-1].content,
        )
        system_prompt = SYSTEM_PROMPT_USELLM_FOR_UNKNOWN.format(role_prompt=role_prompt)
    return Prompt(system_prompt=system_prompt, user_prompt=user_prompt)


async def run_agent(
    knowledges: List[Knowledge], histories: List[Message], role_prompt: str, exclusion: List[Dict] = None
) -> str:
    """Run the agent.

    Args:
        context: The context.
    """
    prompt: Prompt = build_prompt(knowledges, histories, role_prompt, exclusion=exclusion)

    chat_completion = await client.chat.completions.create(
        messages=[
            {"role": "system", "content": prompt.system_prompt},
            {"role": "user", "content": prompt.user_prompt},
        ],
        model=DEFAULT_MODEL_NAME,
    )

    response_str = chat_completion.choices[0].message.content
    return response_str
